import pdfplumber
import fitz  # PyMuPDF
import os
import re
from typing import List, Dict, Tuple, Optional
import uuid


class GuidePDFProcessor:
    def __init__(self, pdf_path: str, output_image_dir: str = "./images"):
        """
        初始化指南PDF处理器
        :param pdf_path: PDF文件路径
        :param output_image_dir: 图片输出目录
        """
        self.pdf_path = pdf_path
        self.output_image_dir = output_image_dir
        self.chapter_pattern = re.compile(r'(\d+\.\s*[\u4e00-\u9fa5\w\s]+)')  # 章节标题模式
        self.section_pattern = re.compile(r'(\d+\.\d+\s*[\u4e00-\u9fa5\w\s]+)')  # 小节标题模式
        self.mark_pattern = re.compile(r'(\◆\s*\d+-\d+至?\s*\◆?\s*\d+-\d+)')  # 来源标记模式

        # 创建图片输出目录
        os.makedirs(output_image_dir, exist_ok=True)

        # 存储处理结果
        self.extracted_content = {
            "text_blocks": [],
            "tables": [],
            "image_related_texts": []
        }

        # 当前章节信息
        self.current_chapter = ""
        self.current_section = ""

    def extract_content(self) -> Dict:
        """提取PDF中的所有内容（文本、表格、图片关联文本）"""
        # 提取文本和表格
        with pdfplumber.open(self.pdf_path) as pdf:
            for page_num, page in enumerate(pdf.pages, 1):
                # 提取页面文本
                self._extract_text(page, page_num)

                # 提取页面表格
                self._extract_tables(page, page_num)

        # 提取图片和关联文本
        self._extract_images_and_related_texts()

        return self.extracted_content

    def _extract_text(self, page, page_num: int):
        """提取页面中的基础文本（章节描述/操作步骤）"""
        text = page.extract_text()
        if not text:
            return

        lines = text.split('\n')
        current_block = []

        for line in lines:
            # 检查是否是章节标题
            chapter_match = self.chapter_pattern.match(line.strip())
            section_match = self.section_pattern.match(line.strip())

            if chapter_match or section_match:
                # 如果有当前文本块，先保存
                if current_block:
                    self._save_text_block('\n'.join(current_block), page_num)
                    current_block = []

                # 更新当前章节/小节
                if chapter_match:
                    self.current_chapter = chapter_match.group(1)
                    self.current_section = ""
                else:
                    self.current_section = section_match.group(1)

                # 将标题本身也作为文本块保存
                self._save_text_block(line.strip(), page_num)
            else:
                # 过滤封面、空白页、重复目录内容
                if not self._is_cover_or_blank(line.strip()):
                    current_block.append(line.strip())

        # 保存最后一个文本块
        if current_block:
            self._save_text_block('\n'.join(current_block), page_num)

    def _save_text_block(self, text: str, page_num: int):
        """保存文本块到提取结果中"""
        # 提取来源标记
        source_marks = self.mark_pattern.findall(text)
        source_mark = ', '.join(source_marks) if source_marks else ""

        # 确定用户角色（这里简单示例，实际可能需要更复杂的逻辑）
        user_role = self._determine_user_role(text)

        self.extracted_content["text_blocks"].append({
            "content": text,
            "page_num": page_num,
            "chapter": self.current_chapter,
            "section": self.current_section,
            "source_mark": source_mark,
            "user_role": user_role
        })

    def _extract_tables(self, page, page_num: int):
        """提取页面中的表格"""
        tables = page.extract_tables()
        if not tables:
            return

        for table in tables:
            # 转换为Markdown格式
            markdown_table = self._convert_to_markdown_table(table)

            # 提取表格前后的文本作为说明
            table_context = self._get_table_context(page, table)

            # 提取来源标记
            source_marks = self.mark_pattern.findall(table_context)
            source_mark = ', '.join(source_marks) if source_marks else ""

            # 确定用户角色
            user_role = self._determine_user_role(table_context)

            self.extracted_content["tables"].append({
                "content": markdown_table,
                "context": table_context,
                "page_num": page_num,
                "chapter": self.current_chapter,
                "section": self.current_section,
                "source_mark": source_mark,
                "user_role": user_role
            })

    def _extract_images_and_related_texts(self):
        """提取图片和关联文本"""
        doc = fitz.open(self.pdf_path)

        for page_num, page in enumerate(doc, 1):
            # 获取页面中的图片
            images = page.get_images(full=True)

            for img_idx, img in enumerate(images):
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]

                # 生成图片文件名（按章节+场景命名）
                chapter_num = self._extract_chapter_number(self.current_chapter)
                image_name = f"{chapter_num}-{self._generate_image_name(page_num, img_idx)}.png"
                image_path = os.path.join(self.output_image_dir, image_name)

                # 保存图片
                with open(image_path, "wb") as f:
                    f.write(image_bytes)

                # 提取图片前后的文本
                img_rect = page.get_image_rects(xref)[0]
                related_text = self._get_image_related_text(doc, page_num, img_rect)

                # 提取来源标记
                source_marks = self.mark_pattern.findall(related_text)
                source_mark = ', '.join(source_marks) if source_marks else ""

                # 确定用户角色
                user_role = self._determine_user_role(related_text)

                self.extracted_content["image_related_texts"].append({
                    "content": related_text,
                    "image_path": image_path,
                    "page_num": page_num,
                    "chapter": self.current_chapter,
                    "section": self.current_section,
                    "source_mark": source_mark,
                    "user_role": user_role
                })

        doc.close()

    def split_text_blocks(self) -> List[Dict]:
        """
        按照规则切割文本块
        遵循"章节分层切割+特殊内容整体保留"原则
        """
        split_blocks = []

        # 处理基础文本块
        for text_block in self.extracted_content["text_blocks"]:
            blocks = self._split_basic_text(text_block)
            split_blocks.extend(blocks)

        # 处理表格
        for table in self.extracted_content["tables"]:
            blocks = self._split_table(table)
            split_blocks.extend(blocks)

        # 处理图片关联文本
        for img_text in self.extracted_content["image_related_texts"]:
            blocks = self._split_image_related_text(img_text)
            split_blocks.extend(blocks)

        return split_blocks

    def _split_basic_text(self, text_block: Dict) -> List[Dict]:
        """切割基础文本（章节描述/操作步骤）"""
        text = text_block["content"]
        blocks = []

        # 如果文本长度小于等于500字符，直接作为一个块
        if len(text) <= 500:
            blocks.append({**text_block, "content": text})
            return blocks

        # 按"章节>子>主题"分层切割
        # 首先尝试按标点符号分割
        sentences = re.split(r'([。！？；,.!?;])', text)
        combined_sentences = []

        # 重新组合句子和标点
        for i in range(0, len(sentences), 2):
            if i + 1 < len(sentences):
                combined_sentences.append(sentences[i] + sentences[i + 1])
            else:
                combined_sentences.append(sentences[i])

        current_block = ""
        for sentence in combined_sentences:
            # 检查添加当前句子后是否超过500字符
            if len(current_block) + len(sentence) <= 500:
                current_block += sentence
            else:
                # 保存当前块
                if current_block:
                    blocks.append({**text_block, "content": current_block})

                # 开始新块，保留前50字符重叠
                overlap = min(50, len(sentence))
                current_block = sentence[-overlap:] if len(current_block) > overlap else sentence

        # 添加最后一个块
        if current_block:
            blocks.append({**text_block, "content": current_block})

        return blocks

    def _split_table(self, table: Dict) -> List[Dict]:
        """切割表格内容"""
        # 表格整体作为1个文本块，除非长度超过500字符
        full_content = f"表格：{table['context']}\n{table['content']}"

        if len(full_content) <= 500:
            return [{**table, "content": full_content}]

        # 如果超过500字符，按"表头+前3行"+"后续行+总结"拆分
        lines = table["content"].split('\n')
        header = []
        body = []
        in_header = True

        for line in lines:
            if in_header and (line.startswith('|') and '-' in line):
                header.append(line)
                in_header = False
            elif in_header:
                header.append(line)
            else:
                body.append(line)

        # 表头部分
        header_content = f"表格：{table['context']}\n" + '\n'.join(header)
        blocks = [{**table, "content": header_content}]

        # 处理表体
        current_block = ""
        for i, line in enumerate(body):
            if i < 3:  # 前3行
                if not current_block:
                    current_block = header_content + '\n'
                current_block += line + '\n'
            else:
                if len(current_block) + len(line) <= 500:
                    current_block += line + '\n'
                else:
                    blocks.append({**table, "content": current_block.strip()})
                    current_block = line + '\n'

        # 添加最后一个块
        if current_block:
            blocks.append({**table, "content": current_block.strip()})

        return blocks

    def _split_image_related_text(self, img_text: Dict) -> List[Dict]:
        """切割图片关联文本"""
        full_content = f"图片说明：{img_text['content']}\n图片路径：{img_text['image_path']}"

        if len(full_content) <= 500:
            return [{**img_text, "content": full_content}]

        # 按"模块"拆分
        modules = re.split(r'(操作步骤|界面模块|功能说明)', full_content)
        blocks = []

        current_block = ""
        for module in modules:
            if not module:
                continue

            if len(current_block) + len(module) <= 500:
                current_block += module
            else:
                if current_block:
                    blocks.append({**img_text, "content": current_block})
                current_block = module

        if current_block:
            blocks.append({**img_text, "content": current_block})

        return blocks

    def add_metadata(self, split_blocks: List[Dict]) -> List[Dict]:
        """为切割后的文本块添加元数据"""
        result = []

        for block in split_blocks:
            # 生成唯一标识
            chapter_num = self._extract_chapter_number(block["chapter"])
            block_id = f"boss_guide_{chapter_num}_{str(uuid.uuid4())[:8]}"

            # 确定内容类型
            content_type = self._determine_content_type(block)

            # 构建元数据
            metadata = {
                "text_block_id": block_id,
                "user_role": block["user_role"],
                "chapter": block["chapter"],
                "content_type": content_type,
                "source_mark": block["source_mark"],
                "pdf_page": block["page_num"],
                "image_path": block.get("image_path", "")
            }

            # 添加元数据到块
            result.append({
                "content": block["content"],
                "metadata": metadata
            })

        return result

    def process(self) -> List[Dict]:
        """完整处理流程：提取内容 -> 切割文本 -> 添加元数据"""
        self.extract_content()
        split_blocks = self.split_text_blocks()
        return self.add_metadata(split_blocks)

    # 辅助方法
    def _is_cover_or_blank(self, text: str) -> bool:
        """判断是否是封面、空白页或重复目录内容"""
        if not text.strip():
            return True

        # 可以添加更多判断规则
        return False

    def _convert_to_markdown_table(self, table: List[List[str]]) -> str:
        """将提取的表格转换为Markdown格式"""
        if not table or len(table) < 2:
            return ""

        markdown = []
        # 添加表头
        markdown.append('| ' + ' | '.join([cell or '' for cell in table[0]]) + ' |')
        # 添加分隔线
        markdown.append('| ' + ' | '.join(['---'] * len(table[0])) + ' |')
        # 添加内容行
        for row in table[1:]:
            markdown.append('| ' + ' | '.join([cell or '' for cell in row]) + ' |')

        return '\n'.join(markdown)

    def _get_table_context(self, page, table) -> str:
        """获取表格前后的文本作为上下文"""
        # 这里简化实现，实际应根据表格位置提取前后文本
        text = page.extract_text() or ""
        return text[:200] + "..." + text[-200:]  # 取表格前后各200字符

    def _generate_image_name(self, page_num: int, img_idx: int) -> str:
        """生成图片文件名"""
        # 这里可以根据实际内容生成更有意义的文件名
        return f"page{page_num}_img{img_idx}"

    def _extract_chapter_number(self, chapter: str) -> str:
        """从章节标题中提取章节号"""
        match = re.match(r'(\d+\.\d*)', chapter)
        if match:
            return match.group(1)
        return f"unknown_{str(uuid.uuid4())[:4]}"

    def _get_image_related_text(self, doc, page_num: int, img_rect: Tuple) -> str:
        """获取图片相关的文本（标题+操作说明）"""
        # 获取当前页文本
        page = doc[page_num - 1]
        text = page.get_text() or ""

        # 这里简化实现，实际应根据图片位置提取相关文本
        return text[:300] + (text[300:] and "...")  # 取前300字符作为相关文本

    def _determine_user_role(self, text: str) -> str:
        """确定用户角色（企业HR/求职者/通用）"""
        # 简单规则示例，实际可能需要更复杂的逻辑
        if "企业" in text or "HR" in text or "招聘" in text:
            return "企业HR"
        elif "求职" in text or "简历" in text or "应聘者" in text:
            return "求职者"
        else:
            return "通用"

    def _determine_content_type(self, block: Dict) -> str:
        """确定内容类型（基础文本/表格/图片关联文本）"""
        if "image_path" in block and block["image_path"]:
            return "图片关联文本"
        elif "content" in block and block["content"].startswith("表格："):
            return "表格"
        else:
            return "基础文本"


# 使用示例
if __name__ == "__main__":
    # 初始化处理器
    pdf_processor = GuidePDFProcessor("Boss 直聘平台使用指南.pdf")

    # 处理PDF
    processed_blocks = pdf_processor.process()

    # 输出结果
    print(f"处理完成，共生成 {len(processed_blocks)} 个带元数据的文本块")

    # 打印前3个块作为示例
    for i, block in enumerate(processed_blocks[:3]):
        print(f"\n文本块 {i + 1}:")
        print(f"内容: {block['content'][:100]}...")  # 只显示前100字符
        print(f"元数据: {block['metadata']}")
